# -*- coding: utf-8 -*-
# @Time    : 2021/4/19
# @File    : the_models.py
from collections import namedtuple

import torch.nn as nn
import torch.nn.functional as F
from torch.nn import Module, Sequential, ReflectionPad2d, InstanceNorm2d, Conv2d
from torchvision import models


class VGG16(nn.Module):
	def __init__(self, requires_grad=False):
		super(VGG16, self).__init__()
		vgg_pretrained_features = models.vgg16(pretrained=True).features
		self.slice1 = nn.Sequential()
		self.slice2 = nn.Sequential()
		self.slice3 = nn.Sequential()
		self.slice4 = nn.Sequential()
		for x in range(4): self.slice1.add_module(str(x), vgg_pretrained_features[x])
		for x in range(4, 9): self.slice2.add_module(str(x), vgg_pretrained_features[x])
		for x in range(9, 16): self.slice3.add_module(str(x), vgg_pretrained_features[x])
		for x in range(16, 23): self.slice4.add_module(str(x), vgg_pretrained_features[x])
		if not requires_grad:
			for param in self.parameters():
				param.requires_grad = False

	def forward(self, x):
		h = self.slice1(x)
		h_relu1_2 = h
		h = self.slice2(h)
		h_relu2_2 = h
		h = self.slice3(h)
		h_relu3_3 = h
		h = self.slice4(h)
		h_relu4_3 = h
		vgg_outputs = namedtuple("VggOutputs", ["relu1_2", "relu2_2", "relu3_3", "relu4_3"])
		out = vgg_outputs(h_relu1_2, h_relu2_2, h_relu3_3, h_relu4_3)
		return out


class ConvBlock(Module):
	def __init__(self, in_channels, out_channels, kernel_size, stride=1, upsample=False, normalize=True, relu=True):
		super(ConvBlock, self).__init__()
		self.upsample = upsample
		self.block = Sequential(
			ReflectionPad2d(kernel_size // 2),
			Conv2d(in_channels, out_channels, kernel_size, stride)
		)
		self.norm = InstanceNorm2d(out_channels, affine=True) if normalize else None
		self.relu = relu

	def forward(self, x):
		if self.upsample: x = F.interpolate(x, scale_factor=2)
		x = self.block(x)
		if self.norm is not None: x = self.norm(x)
		if self.relu: x = F.relu(x)
		return x


# https://arxiv.org/pdf/1512.03385.pdf
# https://zhuanlan.zhihu.com/p/28413039
class ResidualBlock(Module):
	def __init__(self, channels):
		super(ResidualBlock, self).__init__()
		self.block = Sequential(
			ConvBlock(in_channels=channels,
					  out_channels=channels,
					  kernel_size=3,
					  stride=1,
					  normalize=True,
					  relu=True),
			ConvBlock(in_channels=channels,
					  out_channels=channels,
					  kernel_size=3,
					  stride=1,
					  normalize=True,
					  relu=False),
		)

	def forward(self, x):
		return self.block(x) + x


class TransformerNet(Module):
	def __init__(self):
		super(TransformerNet, self).__init__()
		self.model = Sequential(
			ConvBlock(in_channels=3, out_channels=32, kernel_size=9, stride=1),
			ConvBlock(in_channels=32, out_channels=64, kernel_size=3, stride=2),
			ConvBlock(in_channels=64, out_channels=128, kernel_size=3, stride=2),
			ResidualBlock(128),
			ResidualBlock(128),
			ResidualBlock(128),
			ResidualBlock(128),
			ResidualBlock(128),
			ConvBlock(in_channels=128, out_channels=64, kernel_size=3, upsample=True),
			ConvBlock(in_channels=64, out_channels=32, kernel_size=3, upsample=True),
			ConvBlock(in_channels=32, out_channels=3, kernel_size=9,
					  stride=1, normalize=False, relu=False),
		)

	def forward(self, x):
		return self.model(x)
